import os
import numpy as np
from tqdm import tqdm
import cv2
import lazypredict
from lazypredict.Supervised import LazyClassifier
from sklearn.model_selection import train_test_split
import pickle

# 配置数据集路径
data_dir = 'data/train'  # 请替换为你的实际数据集路径

# 定义图像大小
img_size = (64, 64)

# 加载数据集
def load_data(data_dir):
    data = []
    labels = []
    for img in tqdm(os.listdir(data_dir), desc='加载图像'):
        try:
            img_path = os.path.join(data_dir, img)
            img_array = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
            resized_array = cv2.resize(img_array, img_size)
            data.append(resized_array.flatten())
            # 根据文件名提取标签
            if 'cat' in img.lower():
                labels.append(0)
            elif 'dog' in img.lower():
                labels.append(1)
            else:
                continue
        except Exception as e:
            pass
    return np.array(data), np.array(labels)

# 加载数据
X, y = load_data(data_dir)
print(f"数据集大小：{X.shape}")
print(f"标签集大小：{y.shape}")

# 分割数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 使用LazyClassifier进行训练
clf = LazyClassifier(verbose=0, ignore_warnings=True, custom_metric=None)
models, predictions = clf.fit(X_train, X_test, y_train, y_test)

# 输出模型性能
print(models)
# 打印指定的结果列
results = models[['Balanced Accuracy', 'ROC AUC', 'F1 Score', 'Time Taken']]
print(results)

# 获取准确率最高的模型
best_model_name = models.index[0]
best_model = clf.models[best_model_name]

# 保存最佳模型
with open('best_model.pkl', 'wb') as f:
    pickle.dump(best_model, f)

print(f"最佳模型: {best_model_name}")